from selenium import webdriver
from selenium.webdriver.firefox.service import Service as FirefoxService
from webdriver_manager.firefox import GeckoDriverManager


# 优化多浏览器版本
def init_driver(browser="chrome", headless=False):
    '''
    func: 初始化浏览器，功能支持Chrome，Firefox，Safari
    :param browser: 默认chrome
    :param headless:
    :return:
    '''
    if browser.lower() == "chrome":

        options = webdriver.ChromeOptions()
        options.add_argument("--no-sandbox")
        options.add_argument("--disable-dev-shm-usage")
        if headless:
            options.add_argument("--headless")
        return webdriver.Chrome(options=options)

    # 下面的Firefox、Safari没有多浏览器需要可以注释或删除
    elif browser.lower() == "firefox":
        options = webdriver.FirefoxOptions()
        if headless:
            options.add_argument("--headless")
        return webdriver.Firefox(
            service=FirefoxService(GeckoDriverManager().install()),
            options=options
        )

    elif browser == "safari":
        # Safari 没有 options 概念，也不支持 headless
        if headless:
            raise ValueError("Safari does not support headless mode.")
        return webdriver.Safari()

    else:
        raise ValueError(f"Unsupported browser: {browser}")
